package org.dromara.milvus.plus.processor;

import com.google.auto.service.AutoService;
import com.squareup.javapoet.*;
import org.dromara.milvus.plus.annotation.MilvusField;

import javax.annotation.processing.*;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.*;
import javax.lang.model.util.Elements;
import javax.tools.Diagnostic;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;

@SupportedAnnotationTypes("org.dromara.milvus.plus.annotation.GenerateMilvusMapper")
@SupportedSourceVersion(SourceVersion.RELEASE_8)
@AutoService(Processor.class)
public class GenerateMilvusMapperProcessor extends AbstractProcessor {

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        for (TypeElement annotation : annotations) {
            for (Element element : roundEnv.getElementsAnnotatedWith(annotation)) {
                if (element.getKind() == ElementKind.CLASS) {
                    TypeElement typeElement = (TypeElement) element;
                    generateMilvusMapperClass(typeElement);
                    generateConstantsClass(typeElement);
                }
            }
        }
        return true;
    }

    private void generateMilvusMapperClass(TypeElement typeElement) {
        // 获取元素实用工具和类型实用工具
        Elements elementUtils = processingEnv.getElementUtils();
        // 获取实体类的包名和类名
        String packageName = elementUtils.getPackageOf(typeElement).getQualifiedName().toString();
        String className = typeElement.getSimpleName().toString();
        String mapperClassName = className + "MilvusMapper";
        // 创建泛型参数
        TypeName entityTypeName = ClassName.get(packageName, className);
        // 创建Mapper类的TypeSpec
        TypeSpec mapperClass = TypeSpec.classBuilder(mapperClassName)
                .addModifiers(Modifier.PUBLIC)
                .superclass(ParameterizedTypeName.get(
                        ClassName.get("org.dromara.milvus.plus.mapper", "MilvusMapper"),
                        entityTypeName))
                .addAnnotation(AnnotationSpec.builder(ClassName.get("org.springframework.stereotype", "Component")).build())
                .build();

        // 创建JavaFile
        JavaFile javaFile = JavaFile.builder(packageName, mapperClass)
                .addFileComment("Generated by @GenerateMilvusMapper")
                .build();

        try {
            // 写入文件
            javaFile.writeTo(processingEnv.getFiler());
        } catch (IOException e) {
            processingEnv.getMessager().printMessage(Diagnostic.Kind.ERROR, "Error generating Milvus mapper class: " + e.getMessage(), typeElement);
        }
    }
    private void generateConstantsClass(TypeElement typeElement) {
        Elements elementUtils = processingEnv.getElementUtils();
        String packageName = elementUtils.getPackageOf(typeElement).getQualifiedName().toString();
        String className = typeElement.getSimpleName().toString();
        String constantsClassName = className + "Constants";
        List<FieldSpec> constantsFields = new ArrayList<>();
        // 遍历实体类的所有字段
        for (Element enclosedElement : typeElement.getEnclosedElements()) {
            if (enclosedElement.getKind() == ElementKind.FIELD) {
                VariableElement field = (VariableElement) enclosedElement;
                // 检查字段是否有 MilvusField 注解
                MilvusField milvusField = field.getAnnotation(MilvusField.class);
                if (milvusField != null) {
                    String constantName = !milvusField.name().isEmpty() ? milvusField.name() : field.getSimpleName().toString();
                    // 确保常量名称是大写的
                    String upperCaseConstantName = constantName.toUpperCase();
                    constantsFields.add(FieldSpec.builder(String.class, upperCaseConstantName, Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL)
                            .initializer("$S", constantName)
                            .build());
                }
            }
        }
        // 如果有注解字段，则创建 Constants 类
        if (!constantsFields.isEmpty()) {
            TypeSpec constantsClass = TypeSpec.classBuilder(constantsClassName)
                    .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                    .addFields(constantsFields)
                    .build();

            JavaFile javaFile = JavaFile.builder(packageName, constantsClass)
                    .addFileComment("Generated by @GenerateMilvusMapperProcessor")
                    .build();

            try {
                javaFile.writeTo(processingEnv.getFiler());
            } catch (IOException e) {
                processingEnv.getMessager().printMessage(Diagnostic.Kind.ERROR, "Error generating constants class: " + e.getMessage(), typeElement);
            }
        }
    }
}